import logging
import torch
from fusion_vrp.utils.available_window import edge_mask_from_ids

from fusion_vrp.utils.general import cosine_similarity, merge_nodes_v2

__all__ = ['fusion']


def fusion(g, stop_num=5, dist_threshold=1.0):
    g1 = g.clone()

    def shorten_edge(edges):
        dist = edges.data['dist']
        src_feat = edges.src['feat']
        dst_feat = edges.dst['feat']
        # 计算车辆负载比（显式转为 float，保证后续与相似度计算类型一致）
        goods_norm = (edges.src['goods'].float() + edges.dst['goods'].float()) / 100.0
        # 计算时间窗口差


        sim = cosine_similarity(src_feat, dst_feat)
        # Masks by goods_norm ranges
        over_mask = goods_norm > 1.0
        high_mask = (goods_norm > 0.95) & (goods_norm <= 1.0)
        mid_mask = (goods_norm > 0.5) & (goods_norm <= 0.95)
        low_mask = goods_norm <= 0.5

        # Update rules
        # dist_penalty = dist + 100.0                 # overload -> enlarge
        dist_shrink_sim = dist / torch.exp(sim)     # default shrink by similarity
        dist_shrink_mid = dist * 0.5                # gentle shrink for (0.5, 0.9)
        dist_shrink_over = dist + 10
        dist_shrink_low = dist * 0.9                # stronger shrink for <= 0.5
        dist_shrink_high = dist * 0.5


        # Compose by priority: overload > mid/low rules > similarity
        dist_new = dist_shrink_sim
        dist_new = torch.where(mid_mask, dist_shrink_mid, dist_new)
        dist_new = torch.where(high_mask, dist_shrink_high, dist_new)
        dist_new = torch.where(low_mask, dist_shrink_low, dist_new)
        dist_new = torch.where(over_mask, dist_shrink_over, dist_new)


        # 用时间窗口调整距离
        dist_window = dist_new


        logging.debug('旧距离： %s', dist)
        logging.debug('容量比： %s', goods_norm)
        logging.debug('相似度exp： %s', torch.exp(sim))
        logging.debug('新距离： %s', dist_new)
        return {'dist': dist_new}

    count = 1
    while g.num_nodes() > stop_num:
        g.apply_edges(shorten_edge)
        # dists = g.edata['dist']

        # ====================start=================================================
        # 1. 删除所有负载大于1的边（goods特征的和>100）
        src_nodes, dst_nodes, eids = g.edges(form='all')
        goods = g.ndata['goods']
        goods_sum = goods[src_nodes] + goods[dst_nodes]
        remove_mask = goods_sum > 100
        remove_count = int(remove_mask.sum().item())
        if remove_count > 0:
            g.remove_edges(eids[remove_mask])

        # 2. 找到available_window完全相同的节点对，并标记出来
        aw = g.ndata['available_window']
        if aw.dim() > 1 and aw.size(-1) == 1:
            aw = aw.squeeze(-1)
        # 删除所有连接时间窗完全相同的不同节点之间的边（包括双向）
        src_nodes_all, dst_nodes_all, eids_all = g.edges(form='all')
        # mask：节点不同且时间窗相同（使用显式 torch API）
        # ========================================================
        mask = edge_mask_from_ids(g=g, g1=g1, aw1=None, id_key='id', available_window_key='available_window',
                                  debug=False)
        src, dst, eids = g.edges(form='all')
        eids_to_remove = eids[mask]
        if eids_to_remove.numel() > 0:
            logging.debug('删除的边： %s', eids_to_remove)
            g.remove_edges(eids_to_remove)
        # ========================================================
        # eq对相同形状的张量不做广播，逐元素比较，true表示要删除的边

        # 在删除边后，重新读取 dists 以保证后续通过 min_eid 索引边时一致
        dists = g.edata['dist']
        if dists.numel() == 0:
            logging.warning('No edges left after pruning; stopping fusion loop.')
            break
        # 你可以用 same_aw 或 same_aw_triu 进一步处理或记录
        # logging.debug('same available_window pairs: %s', torch.nonzero(same_aw_triu, as_tuple=False))
        # ====================end=================================================
        min_dist, min_eid = torch.min(dists, dim=0)
        if min_dist.item() < dist_threshold:
            src_nodes, dst_nodes = g.edges()
            src = src_nodes[min_eid].item()
            dst = dst_nodes[min_eid].item()
            g = merge_nodes_v2(g, src, dst)
        else:
            # 无效合并次数受限
            count += 1
        if count > 1000:
            logging.error('Infinite loop detected')
            break

    return g
